
import sys
import cv2


# =========================================
def get_selective_search():
    """
    获取selective search算法
    :return:
    """
    return cv2.ximgproc.segmentation.createSelectiveSearchSegmentation()


# =========================================

def config(gs, img, strategy='q'):
    """
    配置selective search 的候选框生成模式
    :param gs: cv2 selective search 算法
    :param img: 图像
    :param strategy: 模式
    :return:
    """
    gs.setBaseImage(img)
    if strategy == 's':
        gs.switchToSingleStrategy()
    if strategy == 'f' :
        gs.switchToSelectiveSearchFast()
    if strategy == 'q':
        gs.switchToSelectiveSearchQuality()
    else:
        print(__doc__)
        sys.exit(1)
# =========================================


def get_rects(gs):
    """
    候选框处理
    :param gs:
    :return:
    """
    try:
        rects = gs.process()
    except Exception as e:
        return []
    rects[:, 2] += rects[:, 0]
    rects[:, 3] += rects[:, 1]

    return rects
# ===============测试==========================


def test(img, color, rects):
    # 循环 region proposal里面的20个候选框
    for x1, y1, x2, y2 in rects[0:20]:
        cv2.rectangle(img, (x1, y1), (x2, y2), color, thickness=2)
    cv2.imshow('head', img)
    cv2.waitKey()


if __name__ == '__main__':

    ss = get_selective_search()

    img = cv2.imread(r'C:\Users\wang1\PycharmProjects\R-CNN-Demo\data\VOC_dog\train\JPEGImages\000036.jpg', cv2.IMREAD_COLOR)
    config(ss, img, strategy='q')
    # 得到region proposal
    rect = get_rects(ss)

    test(img, (255, 255, 255), rect)
